
#ifndef AFTERMATH_PROBABILITY_DIST_LOGNORMAL_HPP_INCLUDED
#define AFTERMATH_PROBABILITY_DIST_LOGNORMAL_HPP_INCLUDED

#include "../double_constants.hpp"
#include "dist_traits.hpp"

#include <cmath>
#include <cstdint>
#include <string>
#include <utility>

namespace aftermath
{
    namespace probability
    {
        // Forward declaration to define traits.
        struct dist_lognormal;

        template <>
        struct is_continuous<dist_lognormal>
        {
            typedef dist_lognormal distribution_type;

            static const bool value = true;
        };

        struct dist_lognormal
        {
            typedef dist_lognormal type;
            typedef double result_type;

            static const std::string name;

        private:
            double m_mu, m_sigma;
            double m_cache_expected_value, m_cache_variance, m_cache_standard_deviation, m_cache_sigma_squared;
            double m_cache_sigma_root_two, m_cache_pdf_scale;

        public:
            dist_lognormal()
                : m_mu(0.0), m_sigma(1.0),
                m_cache_expected_value(std::exp(1.0 / 2)),
                m_cache_variance((std::exp(1.0) - 1.0) * std::exp(1.0)),
                m_cache_standard_deviation(std::sqrt(std::exp(1.0) - 1.0) * std::exp(0.5)),
                m_cache_sigma_squared(1.0),
                m_cache_sigma_root_two(double_constants::root_two), // std::sqrt(2)
                m_cache_pdf_scale(double_constants::one_div_root_two_pi) // 1 / std::sqrt(2 pi)
            {
            }

            explicit dist_lognormal(double mu, double sigma)
                : m_mu(mu), m_sigma(sigma),
                m_cache_expected_value(std::exp(mu + sigma * sigma / 2)),
                m_cache_variance((std::exp(sigma * sigma) - 1.0) * std::exp(2 * mu + sigma * sigma)),
                m_cache_standard_deviation(std::sqrt(std::exp(sigma * sigma) - 1.0) * std::exp(mu + sigma * sigma / 2)),
                m_cache_sigma_squared(sigma * sigma),
                m_cache_sigma_root_two(sigma * double_constants::root_two), // sigma * std::sqrt(2)
                m_cache_pdf_scale(double_constants::one_div_root_two_pi / sigma) // 1 / (sigma * std::sqrt(2 pi))
            {
            }

            double expected_value() const { return this->m_cache_expected_value; }
            double variance() const { return this->m_cache_variance; }
            double standard_deviation() const { return this->m_cache_standard_deviation; }

            double mean() const { return this->expected_value(); }
            double stddev() const { return this->standard_deviation(); }

            double get_mu() const { return this->m_mu; }
            double get_sigma() const { return this->m_sigma; }

            double cdf(double x) const { return std::erfc((this->m_mu - std::log(x)) / (this->m_cache_sigma_root_two)) / 2; }
            double pdf(double x) const { return (this->m_cache_pdf_scale / x) * std::exp(-(std::log(x) - this->m_mu) * (std::log(x) - this->m_mu) / (2 * this->m_cache_variance)); }

            template <std::size_t t_nth_moment>
            double partial_moment(double a, double b) const
            {
                return std::exp(t_nth_moment * this->m_mu + t_nth_moment * t_nth_moment * this->m_cache_sigma_squared / 2) * (
                    std::erf((this->m_mu + t_nth_moment * this->m_cache_sigma_squared - std::log(a)) / (this->m_cache_sigma_root_two)) -
                    std::erf((this->m_mu + t_nth_moment * this->m_cache_sigma_squared - std::log(b)) / (this->m_cache_sigma_root_two))) / 2;
            }

            bool operator ==(const type& other) const
            {
                return
                    this->m_mu == other.m_mu &&
                    this->m_sigma == other.m_sigma;
            }

            bool operator !=(const type& other) const
            {
                return !this->operator==(other);
            }
        };
        
        const std::string dist_lognormal::name = "lognorm";
    }
}

namespace std
{
    template <>
    struct hash<aftermath::probability::dist_lognormal>
    {
        typedef aftermath::probability::dist_lognormal argument_type;
        typedef std::size_t result_type;

        result_type operator()(argument_type const& x) const
        {
            return std::hash<double>()(x.get_mu()) ^ std::hash<double>()(x.get_sigma());
        }
    };
}

#endif // AFTERMATH_PROBABILITY_DIST_LOGNORMAL_HPP_INCLUDED
